from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt 
import sys
import os 
import numpy as np
import time
import threading

X = []
Y = []
Z = []

def startShow():
	dirs = sys.argv[1]
	totalD = os.listdir(dirs)
	keyDict = {}
	for f in totalD:
		label = int(f.split("_")[0])
		if not label in keyDict.keys():
			keyDict[label] = 1
		else:
			keyDict[label] += 1
		if keyDict[label] >= 3:
			continue
		createGraph(dirs+f)

def close(secs):
	time.sleep(secs)
	plt.close()

def createGraph(fname):
	f = np.loadtxt(fname)
	for i in range(len(f)):
		data = f[i]
		X.append(int(data[0]))
		Y.append(int(data[1]))
		Z.append(int(data[2]))
	ax = plt.subplot(111,projection='3d')
	ax.plot_wireframe(X,Y,Z,rstride=1,cstride=1,cmap=plt.cm.Blues_r)
	ax.set_xlabel("X")
	ax.set_ylabel("Y")
	ax.set_zlabel("Z")
	plt.title(fname)
	plt.tight_layout()
	plt.show()		

def getData(dirs,*limit):
	res = []

	for item in dirs:
		label = item.split("_")[0]
		num = (item.split("_")[1]).split(".")[0]
		print(label,int(label) in limit)
		if int(label) not in limit:
			continue
		res.append(item)
		if len(res) >= 20:
			break
	return res

def createBatchGraph():
	dirname = sys.argv[1]
	dirs = os.listdir(dirname)
	dirs = getData(dirs,5)
	fig,axes = plt.subplots(5,4,subplot_kw={'projection':'3d'})
	for i in range(20):
		f = np.loadtxt(dirname+dirs[i])
		row = int(i / 4)
		col = int(i - (row * 4))
		for j in range(len(f)):
			data = f[j]
			X.append(int(data[0]))
			Y.append(int(data[1]))
			Z.append(int(data[2]))
		ax = axes[row][col]
		ax.set_title(dirs[i])
		ax.plot_wireframe(X,Y,Z,rstride=1,cstride=1,cmap=plt.cm.Blues_r)
		ax.set_xlabel("X")
		ax.set_ylabel("Y")
		ax.set_zlabel("Z")
	#plt.title(fname)
	plt.tight_layout()
	plt.show()		

if __name__=='__main__':
	createBatchGraph()

'''
import numpy as np  
import matplotlib.pyplot as plt  
import matplotlib.animation as animation  
  
start = [1, 0.18, 0.63, 0.29, 0.03, 0.24, 0.86, 0.07, 0.58, 0]  
  
metric =[[0.03, 0.86, 0.65, 0.34, 0.34, 0.02, 0.22, 0.74, 0.66, 0.65],  
         [0.43, 0.18, 0.63, 0.29, 0.03, 0.24, 0.86, 0.07, 0.58, 0.55],  
         [0.66, 0.75, 0.01, 0.94, 0.72, 0.77, 0.20, 0.66, 0.81, 0.52]  
        ]  

fig = plt.figure()  
window = fig.add_subplot(111)
help(window)
line, = window.plot(start)  
#如果是参数是list,则默认每次取list中的一个元素,即metric[0],metric[1],...  
def update(data):  
    line.set_ydata(data)  
    return line,  
  
ani = animation.FuncAnimation(fig, update, metric, interval=2*1000)  
plt.show() 

"""
============
3D animation
============

A simple example of an animated plot... In 3D!

import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as p3
import matplotlib.animation as animation
import sys
import os 

def Get_Fname():
	dirs = os.listdir("blue_data/")
	for f in dirs:
		yield "blue_data/" + f

def Gen_RandLine(length, dims=2):
    """"""
    Create a line using a random walk algorithm

    length is the number of points for the line.
    dims is the number of dimensions the line has."""
    """
    lineData = np.empty((dims, length))
    lineData[:, 0] = np.random.rand(dims)
    for index in range(1, length):
        # scaling the random numbers by 0.1 so
        # movement is small compared to position.
        # subtraction by 0.5 is to change the range to [-0.5, 0.5]
        # to allow a line to move backwards.
        step = ((np.random.rand(dims) - 0.5) * 0.1)
        lineData[:, index] = lineData[:, index - 1] + step

    return lineData


def update_lines(num, dataLines, lines):
    for line, data in zip(lines, dataLines):
        # NOTE: there is no .set_data() for 3 dim data...
        line.set_data(data[0:2, :num])
        line.set_3d_properties(data[2, :num])
        help(lines)
        exit()
    return lines

def update(fname, line):
	f = np.loadtxt(fname)
	X = []
	Y = []
	Z = []
	f = np.loadtxt(fname)
	for i in range(len(f)):
		data = f[i]
		X.append(int(data[0]))
		Y.append(int(data[1]))
		Z.append(int(data[2]))
	line.set_data(X,Y)
	line.set_3d_properties(Z)
	return line

# Attaching 3D axis to the figure
fig = plt.figure()
ax = p3.Axes3D(fig)

# Fifty lines of random 3-D lines
data = [Gen_RandLine(25, 3) for index in range(50)]

# Creating fifty line objects.
# NOTE: Can't pass empty arrays into 3d version of plot()
lines = [ax.plot(dat[0, 0:1], dat[1, 0:1], dat[2, 0:1])[0] for dat in data]

# Setting the axes properties
ax.set_xlim3d([0.0, 3.0])
ax.set_xlabel('X')

ax.set_ylim3d([0.0, 3.0])
ax.set_ylabel('Y')

ax.set_zlim3d([0.0, 3.0])
ax.set_zlabel('Z')

ax.set_title('3D Test')

# Creating the Animation object
line_ani = animation.FuncAnimation(fig, update_lines, 26, fargs=(data, lines),
                                   interval=50, blit=False)
#line_ani = animation.FuncAnimation(fig, update, Get_Fname, fargs=(ax), interval=1000)

plt.show()
'''